import torch
import torch.nn as nn


class AttentionMicroscope(nn.Module):
    """AI警探的智能放大镜：聚焦损伤关键证据"""

    def __init__(self, input_dim=128, hidden_dim=64):
        super().__init__()
        # 注意力机制：自动识别重要特征
        self.attention = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, 1),
            nn.Softmax(dim=1)
        )

    def forward(self, features):
        """前向传播：返回加权后的关键特征"""
        attn_weights = self.attention(features)
        return torch.sum(attn_weights * features, dim=1)